from copy import deepcopy
import math
import random
from chiplet import QPU
from tool import remove_root_node_from_dag, extract_first_layer
from math import log2, ceil
import networkx as nx


# ⭐️ 优化选路
def select_sole_path_for_given_CNOTs(qpu: QPU, cnots: list[list]):
    '''待优化的地方，针对给定的一批CNOT门（不一定是全部的CNOT门），选择唯一的路径\n
    传入一批CNOT门，为这些门请求选择路径\n
    为每个CNOT节点对，计算唯一的路径
    '''
    arr = []
    for cnot in cnots:
        tmp = qpu.get_k_path_with_qubit_id(
            source_qubit=cnot[0], target_qubit=cnot[1], k=1)
        arr.append(tmp[0])
    return arr


# 对应的Highway
class Highway:
    def __init__(self, serve_dag_id, generate_time, live_time, ghz_link, last_serve_time):
        self.serve_dag_id: int = serve_dag_id  # 服务的 CNOT 节点
        self.generate_time: int = generate_time  # 开始生成的时间
        self.live_time: int = live_time  # GHZ态的存续时间
        # 相关的ghz链路，tuple中是物理比特ID
        self.ghz_link: list[tuple[int, int]] = ghz_link
        self.last_serve_time: int = last_serve_time  # 上一次服务的完成时间


def remove_invalid_ghz(remain_GHZ_highway: list[Highway], now_time: int):
    '''去掉失效的GHZ Highway'''
    # 找出失效的 Highway 对象
    invalid_highways = [
        highway for highway in remain_GHZ_highway
        if highway.generate_time + highway.live_time < now_time
    ]
    # 更新 remain_GHZ_highway 列表，保留未过期的对象
    remain_GHZ_highway[:] = [
        highway for highway in remain_GHZ_highway
        if highway.generate_time + highway.live_time >= now_time
    ]


def extend_or_new(qpu: QPU, remain_GHZ_highway: list[Highway], new_cnot_dag_id, final_time_dag_optimal: nx.DiGraph, now_time: int):
    '''判断是搭在现有的GHZ上，还是新建一个GHZ\n
    返回选中的Highway的Index，花费的时间'''

    swap_num = 0
    ghz_length = 0

    # 先清除所有失效的GHZ
    remove_invalid_ghz(remain_GHZ_highway, now_time)

    # print("剩余的GHZ数量：", len(remain_GHZ_highway))

    candidate_index = -1
    earliest_time = 1000000
    cnot = final_time_dag_optimal.nodes[new_cnot_dag_id]["label"]
    source_physics_qubit_id = 0
    target_physics_qubit_id = 0
    extend_path_links = []
    for _key, value in qpu.map_plan.items():
        if value == cnot[0]:
            source_physics_qubit_id = _key
        if value == cnot[1]:
            target_physics_qubit_id = _key
    just_build_a_new_highway_time = 0  # 重新开始构建一个Highway
    brand_new_ghz_path = qpu.get_k_path_with_qubit_id(cnot[0], cnot[1], 1)[0]
    tmp_length = len(brand_new_ghz_path)
    just_build_a_new_highway_time += ceil(
        log2(tmp_length if tmp_length > 0 else 1))
    for physic_qubit_id in brand_new_ghz_path:
        if (qpu.map_plan[physic_qubit_id] != None):
            just_build_a_new_highway_time += 1
    for highway_idx, highway in enumerate(remain_GHZ_highway):
        generated_time = highway.generate_time
        live_time = highway.live_time
        final_can_serve_time = generated_time + live_time
        min_time_1 = 1000  # 记录 source 到Highway的最短时间
        min_time_2 = 1000  # 记录 target 到Highway的最短时间
        min_path_1 = []  # 记录source到Highway的最短路径
        min_path_2 = []  # 记录target到Highway的最短时间路径

        for link in highway.ghz_link:
            a = 0
            b = 0
            a, b = link[0], link[1]  # 对应的物理比特ID
            source_to_cnot_start_time, source_to_cnot_end_time, target_to_cnot_start_time, target_to_cnot_end_time = 0, 0, 0, 0

            # 🤔 源节点到link端点a
            tmp_1_1 = nx.shortest_path(
                qpu.merged_graph, source_physics_qubit_id, a)
            tmp_length = len(tmp_1_1) - 1
            source_to_cnot_start_time = ceil(
                log2(tmp_length if tmp_length > 0 else 1))
            tmp_1_1_flag = False
            for i in tmp_1_1:
                if (qpu.map_plan[i] != None):
                    tmp_1_1_flag = True
            if (tmp_1_1_flag):
                source_to_cnot_start_time += 1  # 移动qubit花费的时间
            path_1_1 = [(tmp_1_1[i], tmp_1_1[i+1])
                        for i in range(len(tmp_1_1)-1)]

            # 🤔 源节点到link端点b
            tmp_1_2 = nx.shortest_path(
                qpu.merged_graph, source_physics_qubit_id, b)
            tmp_length = len(tmp_1_2) - 1
            source_to_cnot_end_time = ceil(
                log2(tmp_length if tmp_length > 0 else 1))
            tmp_1_2_flag = False
            for i in tmp_1_2:
                if (qpu.map_plan[i] != None):
                    tmp_1_2_flag = True
            if (tmp_1_2_flag):
                source_to_cnot_end_time += 1  # 移动qubit花费的时间
            path_1_2 = [(tmp_1_2[i], tmp_1_2[i+1])
                        for i in range(len(tmp_1_2)-1)]

            # 🤔 目的节点到link端点a
            tmp_2_1 = nx.shortest_path(
                qpu.merged_graph, target_physics_qubit_id, a)
            tmp_time = len(tmp_2_1) - 1
            target_to_cnot_start_time = ceil(
                log2(tmp_time if tmp_time > 0 else 1))
            tmp_2_1_flag = False
            for i in tmp_2_1:
                if (qpu.map_plan[i] != None):
                    tmp_2_1_flag = True
            if (tmp_2_1_flag):
                target_to_cnot_start_time += 1  # 移动qubit花费的时间
            path_2_1 = [(tmp_2_1[i], tmp_2_1[i+1])
                        for i in range(len(tmp_2_1)-1)]

            # 🤔 目的节点到link端点b
            tmp_2_2 = nx.shortest_path(
                qpu.merged_graph, target_physics_qubit_id, b)
            tmp_time = len(tmp_2_2) - 1
            target_to_cnot_end_time = ceil(
                log2(tmp_time if tmp_time > 0 else 1))
            tmp_2_2_flag = False
            for i in tmp_2_2:
                if (qpu.map_plan[i] != None):
                    tmp_2_2_flag = True
            if (tmp_2_2_flag):
                target_to_cnot_end_time += 1  # 移动qubit花费的时间
            path_2_2 = [(tmp_2_2[i], tmp_2_2[i+1])
                        for i in range(len(tmp_2_2)-1)]

            # 更新最短路径和对应的links
            if source_to_cnot_start_time < min_time_1:
                min_time_1 = source_to_cnot_start_time
                min_path_1 = path_1_1
            if source_to_cnot_end_time < min_time_2:
                min_time_2 = source_to_cnot_end_time
                min_path_2 = path_1_2
            if target_to_cnot_start_time < min_time_1:
                min_time_1 = target_to_cnot_start_time
                min_path_1 = path_2_1
            if target_to_cnot_end_time < min_time_2:
                min_time_2 = target_to_cnot_end_time
                min_path_2 = path_2_2

        min_time = min(min_time_1, min_time_2)  # 最小的搭上时间
        estimate_complete_time = now_time + min_time
        extend_path_links.extend(min_path_1)
        extend_path_links.extend(min_path_2)

        if (estimate_complete_time < final_can_serve_time):
            if (estimate_complete_time < earliest_time):
                earliest_time = estimate_complete_time
                candidate_index = highway_idx

    swap_num = 0
    if (now_time + just_build_a_new_highway_time < earliest_time):
        # ⭐️ 新建一条新的GHZ
        swap_num += get_cross_computing_qubit_num(brand_new_ghz_path, qpu)
        ghz_length += len(brand_new_ghz_path)
        remain_GHZ_highway.append(
            Highway(
                serve_dag_id=new_cnot_dag_id,
                generate_time=now_time,
                live_time=8,  # 存活时间 8 个时隙
                ghz_link=[(brand_new_ghz_path[i], brand_new_ghz_path[i+1])
                          for i in range(len(brand_new_ghz_path)-1)],
                last_serve_time=now_time + just_build_a_new_highway_time
            )
        )
        return -1, just_build_a_new_highway_time, swap_num, ghz_length

    swap_num = 0
    if (candidate_index > -1):
        
        def already_in(physical_id):
            ghz_link = remain_GHZ_highway[candidate_index].ghz_link
            for link in ghz_link:
                if (link[0] == physical_id or link[1] == physical_id):
                    return True
            return False
        
        added_ids = {}
        for _i in extend_path_links:
            _id1 = _i[0]
            _id2 = _i[1]
            if (qpu.map_plan[_id1] != None or qpu.map_plan[_id2] != None):
                if(not already_in(_id1) and not already_in(_id2)):
                    if(qpu.map_plan[_id1] != None):
                        if(_id1 not in added_ids):
                            swap_num += 1
                            added_ids[_id1] = 1
                    if(qpu.map_plan[_id2] != None):
                        if(_id2 not in added_ids):
                            swap_num += 1
                            added_ids[_id2] = 1
        
        ghz_length += len(extend_path_links)
        # ⭐️ 扩展以前旧的GHZ
        remain_GHZ_highway[candidate_index].ghz_link += extend_path_links
        remain_GHZ_highway[candidate_index].last_serve_time = estimate_complete_time

    # 返回最终选中的高速通道
    return candidate_index, estimate_complete_time - now_time, swap_num, ghz_length


def explore_eliminate_ghz(qpu: QPU, cnots_in_window, remain_GHZ_highway):
    '''探索削减GHZ长度的可能性，使用轻量级模拟退火优化'''
    N = 5  # 每次交换选中的Qubit数量
    K = 5  # 每次温度下尝试的解数量
    T = 5  # 温度下降步数，越大越精细但越耗时
    cooling_rate = 0.8  # 冷却速度

    def get_all_consuming(cnot_paths, swap_consume):
        return sum(ceil(log2(len(path))) for path in cnot_paths) + swap_consume

    def can_swap(pq, remain_GHZ_highway: list[Highway]):
        # 不能交换到已有的GHZ上
        flag = True
        for i in remain_GHZ_highway:
            for j in i.ghz_link:
                if (pq == j[0] or pq == j[1]):
                    flag = False
        return flag

    def sample_swap(current_qpu):
        # 随机选择N个计算比特
        logical_qubits = list(current_qpu.map_plan.keys())
        physical_qubits = [
            pq for pq in logical_qubits if current_qpu.map_plan[pq] is not None]

        # 随机选择N个物理量子比特（或少于N个，如果不足）
        selected_pqs = random.sample(
            physical_qubits, min(N, len(physical_qubits)))

        # 记录交换历史
        swap_history = []

        for source_pq in selected_pqs:
            adjacent_pqs = list(current_qpu.merged_graph.neighbors(source_pq))

            # 过滤掉不能交换的邻居
            available_neighbors = [
                pq for pq in adjacent_pqs if can_swap(pq, remain_GHZ_highway)]

            if available_neighbors:
                target_pq = random.choice(available_neighbors)

                # 交换两个物理 Qubit 的映射
                current_qpu.map_plan[source_pq], current_qpu.map_plan[target_pq] = (
                    current_qpu.map_plan[target_pq],
                    current_qpu.map_plan[source_pq]
                )
                swap_history.append((source_pq, target_pq))

        return swap_history

    # 撤销swap操作
    def undo_swaps(current_qpu, swap_history):
        for source_pq, target_pq in reversed(swap_history):
            current_qpu.map_plan[source_pq], current_qpu.map_plan[target_pq] = (
                current_qpu.map_plan[target_pq],
                current_qpu.map_plan[source_pq]
            )

    # 初始化最佳解
    best_solution = deepcopy(qpu)
    best_cost = get_all_consuming(
        select_sole_path_for_given_CNOTs(best_solution, cnots_in_window), 0)

    current_solution = best_solution
    current_cost = best_cost

    temperature = 1.0  # 初始温度

    for _ in range(T):  # 控制温度下降次数
        swap_histories = []
        new_costs = []
        new_solutions = []

        for _ in range(K):  # 尝试K个候选解
            history = sample_swap(current_solution)
            new_cnot_paths = select_sole_path_for_given_CNOTs(
                current_solution, cnots_in_window)
            new_cost = get_all_consuming(new_cnot_paths, len(history))
            swap_histories.append(history)
            new_costs.append(new_cost)
            new_solutions.append(
                deepcopy(current_solution.map_plan))  # 只复制 map_plan
            undo_swaps(current_solution, history)  # 立即撤销本次修改

        # 遍历所有候选解
        for idx in range(K):
            new_cost = new_costs[idx]
            new_map_plan = new_solutions[idx]

            delta = new_cost - current_cost

            # 如果是更优解，或者以一定概率接受劣解
            if delta < 0 or random.random() < math.exp(-delta / temperature):
                current_solution.map_plan = new_map_plan
                current_cost = new_cost

                # 更新全局最优解
                if current_cost < best_cost:
                    best_cost = current_cost
                    best_solution.map_plan = new_map_plan

        temperature *= cooling_rate  # 温度衰减

    return best_solution


def schedule_optimal(original_qpu: QPU, callbackSchedulePath=None):
    '''⭐️ 优化调度，计算最后电路完成时间'''
    qpu: QPU = deepcopy(original_qpu)  # 先拷贝一份
    # ⭐️ 在调度的过程中，量子比特的映射会发生变化
    final_time_dag_optimal = deepcopy(qpu.circuit_dag.dag_nx)

    tmp_dag_optimal = deepcopy(qpu.circuit_dag.dag_nx)
    remove_root_node_from_dag(tmp_dag_optimal)
    has_remaining_cnot = True

    remain_GHZ_highway: list[Highway] = []  # 保留残存的GHZ Highway

    swap_num = 0
    ghz_length = 0

    while has_remaining_cnot:
        # 抽取 Front Layer
        cnots_in_window, tmp_dag_with_attrs, node_id_in_window, _has_remaining_cnot = extract_first_layer(
            tmp_dag_optimal)
        # 判断DAG图中是否有剩余的CNOT门没处理
        has_remaining_cnot = _has_remaining_cnot
        if (not has_remaining_cnot):
            break
        # 采样其中一部分CNOT门
        selected_node_dag_ids = node_id_in_window

        # ⭐️ 探索削减GHZ Length的可能
        qpu = explore_eliminate_ghz(qpu, cnots_in_window, remain_GHZ_highway)

        # ⭐️ 选路
        cnot_paths = select_sole_path_for_given_CNOTs(
            qpu, cnots_in_window)

        # 排序index
        # sorted_indices = sorted(range(len(cnot_paths)),
        #                         key=lambda i: len(cnot_paths[i]))
        # 按照优先级先后顺序排序
        sorted_indices = list(reversed(sorted(range(len(selected_node_dag_ids)),
                                              key=lambda i: final_time_dag_optimal.nodes[selected_node_dag_ids[i]]["time_weight"])))
        # 根据排序后的索引重新排列 selected_node_dag_ids
        sorted_node_dag_ids = [selected_node_dag_ids[i]
                               for i in sorted_indices]
        sorted_paths = [cnot_paths[i] for i in sorted_indices]
        non_conflicting_path_idxs = find_conflict(sorted_paths)
        dag_ids = [sorted_node_dag_ids[i] for i in non_conflicting_path_idxs]
        # print("processed dag：", dag_ids)
        can_schedule_paths = [sorted_paths[i]
                              for i in non_conflicting_path_idxs]
        if (callbackSchedulePath != None):
            callbackSchedulePath(can_schedule_paths, qpu)
        # ⭐️ 要调度的路径
        for idx in non_conflicting_path_idxs:
            dag_node_id = sorted_node_dag_ids[idx]
            path = sorted_paths[idx]
            max_pre_time = 0
            for pre in final_time_dag_optimal.predecessors(dag_node_id):
                max_pre_time = max(
                    max_pre_time, final_time_dag_optimal.nodes[pre]["process_done_time"])
            # 搭上已有的 GHZ 或者创建一条新的链路
            candidate_index, consume_time, swap_num_, ghz_length_ = extend_or_new(
                qpu, remain_GHZ_highway, dag_node_id, final_time_dag_optimal, max_pre_time)

            swap_num += swap_num_
            ghz_length += ghz_length_

            final_time_dag_optimal.nodes[dag_node_id]["process_done_time"] = max_pre_time + consume_time

            diff_idxs = [
                x for x in sorted_indices if x not in non_conflicting_path_idxs]

            for idx1 in diff_idxs:
                dag_node_id1 = sorted_node_dag_ids[idx1]
                path1 = sorted_paths[idx1]
                if (is_two_path_conflict(path, path1)):
                    # 如果他俩冲突的话，就增加一条边，保证DAG的偏序性
                    final_time_dag_optimal.add_edge(
                        dag_node_id, dag_node_id1, weight=1)

        tmp_dag_optimal.remove_nodes_from(
            [sorted_node_dag_ids[i] for i in non_conflicting_path_idxs])

    max_depth = 0
    # 找出 self.final_depth_dag 出度为0的顶点
    for node in final_time_dag_optimal.nodes():
        if (final_time_dag_optimal.nodes[node]["process_done_time"] > max_depth):
            max_depth = final_time_dag_optimal.nodes[node]["process_done_time"]

    return max_depth, swap_num, ghz_length
    print("max_time_optimal: ", max_depth)


def is_two_path_conflict(path1, path2) -> bool:
    '''判断两条路径是否存在冲突'''
    links_set = set()
    for i in range(len(path1) - 1):
        u, v = path1[i], path1[i + 1]
        link = (min(u, v), max(u, v))
        links_set.add(link)
    for i in range(len(path2) - 1):
        u, v = path2[i], path2[i + 1]
        link = (min(u, v), max(u, v))
        if link in links_set:
            return True
    return False


def find_conflict(paths):
    '''找到多条路径彼此不冲突的，或者冲突路径中最短的那一条'''
    link_count = {}
    path_links = []
    for path in paths:
        links = set((min(path[i], path[i+1]), max(path[i], path[i+1]))
                    for i in range(len(path)-1))
        path_links.append(links)
        for link in links:
            link_count[link] = link_count.get(link, 0) + 1
    conflicts = [any(link_count[link] > 1 for link in links)
                 for links in path_links]
    non_conflicting_indices = [
        i for i, conflict in enumerate(conflicts) if not conflict]
    conflicting_indices = [
        i for i, conflict in enumerate(conflicts) if conflict]
    shortest_conflicting_index = min(
        conflicting_indices, key=lambda i: len(paths[i])) if conflicting_indices else None
    return non_conflicting_indices + ([shortest_conflicting_index] if shortest_conflicting_index is not None else [])


def get_cross_computing_qubit_num(path, qpu: QPU):
    '''获取穿过了的量子比特数'''
    count = 0
    for i in path:
        try:
            if (qpu.map_plan[i] != None):
                count += 1
        except Exception as e:
            pass
    return count
